import numpy as np
from enum import Enum, auto

class FlowRegime(Enum):
    """流型枚举类，便于代码维护和扩展"""
    BUBBLY = auto()    # 气泡流
    SLUG = auto()      # 弹状流
    CHURN = auto()     # 搅拌流
    ANNULAR = auto()   # 环状流
    MIST = auto()      # 雾状流
    
    @classmethod
    def to_str(cls, regime):
        """将枚举转换为字符串"""
        return regime.name.lower()

class FlowRegimeIdentifier:
    """增强版流型识别器，基于多参数融合策略"""
    def __init__(self, config=None):
        """
        初始化流型识别器
        
        Args:
            config: 配置参数字典，若为None则使用默认配置
        """
        # 默认配置
        self.config = {
            # 基本阈值配置
            'bubbly_max_alpha': 0.3,
            'slug_max_alpha': 0.5,
            'churn_max_alpha': 0.7,
            
            # 相对速度影响因子
            'velocity_influence': 0.2,
            
            # 管径影响因子
            'diameter_influence': 0.1,
            
            # 流体特性影响因子
            'fluid_property_influence': 0.15,
            
            # 是否启用无量纲参数
            'use_dimensionless_params': True,
            
            # 是否考虑流动方向（水平/垂直）
            'consider_flow_direction': True,
            
            # 默认流动方向（垂直）
            'is_vertical': True
        }
        
        # 更新配置
        if config is not None:
            self.config.update(config)
        
        # 初始化无量纲参数计算器
        self.dimensionless_calculator = DimensionlessParametersCalculator()
    
    def identify_regime(self, alpha, velocity_r, **kwargs):
        """
        基于多参数融合策略识别流型
        
        Args:
            alpha: 空泡率
            velocity_r: 相对速度 (m/s)
            **kwargs: 可选参数
                - diameter: 水力直径 (m)
                - rho_l: 液相密度 (kg/m³)
                - rho_g: 气相密度 (kg/m³)
                - mu_l: 液相动力粘度 (Pa·s)
                - mu_g: 气相动力粘度 (Pa·s)
                - sigma: 表面张力系数 (N/m)
                - is_vertical: 是否垂直流动
                - pressure: 系统压力 (Pa)
                - temperature: 系统温度 (K)
                
        Returns:
            str: 流型类型
        """
        # 提取可选参数
        diameter = kwargs.get('diameter', 0.01)  # 默认1cm
        rho_l = kwargs.get('rho_l', 1000.0)      # 默认水的密度
        rho_g = kwargs.get('rho_g', 1.0)         # 默认空气的密度
        mu_l = kwargs.get('mu_l', 1e-3)          # 默认水的粘度
        mu_g = kwargs.get('mu_g', 1.8e-5)        # 默认空气的粘度
        sigma = kwargs.get('sigma', 0.072)       # 默认水-空气的表面张力
        is_vertical = kwargs.get('is_vertical', self.config['is_vertical'])
        
        # 基础流型判断（基于空泡率）
        base_regime = self._get_base_regime(alpha)
        
        # 如果启用无量纲参数，计算修正因子
        correction_factor = 0.0
        if self.config['use_dimensionless_params']:
            correction_factor = self._calculate_correction_factor(
                alpha, velocity_r, diameter, rho_l, rho_g, mu_l, mu_g, sigma
            )
        
        # 考虑流动方向的影响
        if self.config['consider_flow_direction']:
            if is_vertical:
                # 垂直流动时的修正
                regime = self._adjust_for_vertical_flow(base_regime, alpha, velocity_r, correction_factor)
            else:
                # 水平流动时的修正
                regime = self._adjust_for_horizontal_flow(base_regime, alpha, velocity_r, correction_factor)
        else:
            # 不考虑流动方向，直接应用修正因子
            regime = self._apply_correction(base_regime, correction_factor)
        
        return FlowRegime.to_str(regime)
    
    def _get_base_regime(self, alpha):
        """基于空泡率的基础流型判断"""
        if alpha < self.config['bubbly_max_alpha']:
            return FlowRegime.BUBBLY
        elif alpha < self.config['slug_max_alpha']:
            return FlowRegime.SLUG
        elif alpha < self.config['churn_max_alpha']:
            return FlowRegime.CHURN
        else:
            return FlowRegime.ANNULAR
    
    def _calculate_correction_factor(self, alpha, velocity_r, diameter, rho_l, rho_g, mu_l, mu_g, sigma):
        """计算基于无量纲参数的修正因子"""
        # 计算无量纲参数
        params = self.dimensionless_calculator.calculate_parameters(
            alpha, velocity_r, diameter, rho_l, rho_g, mu_l, mu_g, sigma
        )
        
        # 提取关键无量纲参数
        re_l = params['reynolds_liquid']
        re_g = params['reynolds_gas']
        we = params['weber']
        fr = params['froude']
        
        # 基于无量纲参数计算修正因子
        # 这里使用简化的计算方法，实际应用中可以基于实验数据或更复杂的模型
        correction = 0.0
        
        # Reynolds数的影响（湍流程度）
        if re_l > 2300 and re_g > 2300:
            # 两相都是湍流，更容易形成复杂流型
            correction += 0.1
        
        # Weber数的影响（表面张力vs惯性力）
        if we > 10:
            # 高Weber数，惯性力占主导，更容易破碎
            correction += 0.15
        
        # Froude数的影响（重力vs惯性力）
        if fr < 1:
            # 低Froude数，重力占主导
            correction -= 0.05
        
        return correction
    
    def _adjust_for_vertical_flow(self, base_regime, alpha, velocity_r, correction_factor):
        """针对垂直流动的流型调整"""
        # 垂直流动中，相对速度对流型的影响更大
        velocity_effect = 0.0
        if velocity_r > 0.5:  # 相对速度较大
            if base_regime == FlowRegime.BUBBLY:
                # 高相对速度可能导致气泡流向弹状流转变
                velocity_effect = 0.1
            elif base_regime == FlowRegime.SLUG:
                # 高相对速度可能导致弹状流向搅拌流转变
                velocity_effect = 0.15
        
        # 综合修正因子
        total_correction = correction_factor + velocity_effect
        
        return self._apply_correction(base_regime, total_correction)
    
    def _adjust_for_horizontal_flow(self, base_regime, alpha, velocity_r, correction_factor):
        """针对水平流动的流型调整"""
        # 水平流动中，分层效应更明显
        stratification_effect = 0.0
        if velocity_r < 0.3:  # 相对速度较小
            # 低相对速度下，更容易出现分层流
            stratification_effect = -0.2
        
        # 综合修正因子
        total_correction = correction_factor + stratification_effect
        
        return self._apply_correction(base_regime, total_correction)
    
    def _apply_correction(self, base_regime, correction_factor):
        """应用修正因子到基础流型"""
        # 流型序号
        regime_order = [
            FlowRegime.BUBBLY,
            FlowRegime.SLUG,
            FlowRegime.CHURN,
            FlowRegime.ANNULAR,
            FlowRegime.MIST
        ]
        
        # 获取基础流型的序号
        base_index = regime_order.index(base_regime)
        
        # 应用修正因子
        # 正值表示向更分散的流型转变，负值表示向更聚集的流型转变
        corrected_index = base_index + round(correction_factor * 2)
        
        # 确保索引在有效范围内
        corrected_index = max(0, min(corrected_index, len(regime_order) - 1))
        
        return regime_order[corrected_index]


class DimensionlessParametersCalculator:
    """无量纲参数计算器"""
    
    def calculate_parameters(self, alpha, velocity_r, diameter, rho_l, rho_g, mu_l, mu_g, sigma):
        """
        计算关键无量纲参数
        
        Args:
            alpha: 空泡率
            velocity_r: 相对速度 (m/s)
            diameter: 水力直径 (m)
            rho_l: 液相密度 (kg/m³)
            rho_g: 气相密度 (kg/m³)
            mu_l: 液相动力粘度 (Pa·s)
            mu_g: 气相动力粘度 (Pa·s)
            sigma: 表面张力系数 (N/m)
            
        Returns:
            dict: 包含计算得到的无量纲参数
        """
        # 计算相速度（简化估计）
        velocity_g = velocity_r * 1.5  # 假设气相速度
        velocity_l = velocity_r * 0.5  # 假设液相速度
        
        # 计算Reynolds数
        re_l = self._reynolds_number(rho_l, velocity_l, diameter, mu_l)
        re_g = self._reynolds_number(rho_g, velocity_g, diameter, mu_g)
        
        # 计算Weber数
        we = self._weber_number(rho_l, velocity_r, diameter, sigma)
        
        # 计算Froude数
        fr = self._froude_number(velocity_r, diameter)
        
        # 计算Bond数（Eötvös数）
        bo = self._bond_number(rho_l, rho_g, diameter, sigma)
        
        # 计算Ohnesorge数
        oh = self._ohnesorge_number(rho_l, diameter, mu_l, sigma)
        
        # 计算Capillary数
        ca = self._capillary_number(mu_l, velocity_l, sigma)
        
        return {
            'reynolds_liquid': re_l,
            'reynolds_gas': re_g,
            'weber': we,
            'froude': fr,
            'bond': bo,
            'ohnesorge': oh,
            'capillary': ca
        }
    
    def _reynolds_number(self, rho, velocity, diameter, mu):
        """计算Reynolds数: Re = ρvD/μ"""
        return rho * velocity * diameter / mu
    
    def _weber_number(self, rho, velocity, diameter, sigma):
        """计算Weber数: We = ρv²D/σ"""
        return rho * velocity**2 * diameter / sigma
    
    def _froude_number(self, velocity, diameter):
        """计算Froude数: Fr = v²/(gD)"""
        g = 9.81  # 重力加速度 (m/s²)
        return velocity**2 / (g * diameter)
    
    def _bond_number(self, rho_l, rho_g, diameter, sigma):
        """计算Bond数: Bo = g(ρl-ρg)D²/σ"""
        g = 9.81  # 重力加速度 (m/s²)
        return g * (rho_l - rho_g) * diameter**2 / sigma
    
    def _ohnesorge_number(self, rho, diameter, mu, sigma):
        """计算Ohnesorge数: Oh = μ/√(ρσD)"""
        # 防止除零和负数开方
        if rho <= 0 or sigma <= 0 or diameter <= 0:
            return 0.0
        denominator = np.sqrt(rho * sigma * diameter)
        if denominator <= 1e-15:
            return 0.0
        return mu / denominator
    
    def _capillary_number(self, mu, velocity, sigma):
        """计算Capillary数: Ca = μv/σ"""
        return mu * velocity / sigma


class FlowRegimeMap:
    """流型图实现，用于可视化和验证流型判断"""
    
    def __init__(self, identifier):
        """
        初始化流型图
        
        Args:
            identifier: FlowRegimeIdentifier实例
        """
        self.identifier = identifier
    
    def generate_map(self, x_param, y_param, **fixed_params):
        """
        生成流型图
        
        Args:
            x_param: X轴参数配置 {'name': 参数名, 'range': [最小值, 最大值], 'points': 点数}
            y_param: Y轴参数配置 {'name': 参数名, 'range': [最小值, 最大值], 'points': 点数}
            **fixed_params: 固定参数值
            
        Returns:
            tuple: (X网格, Y网格, 流型矩阵)
        """
        # 创建参数网格
        x_values = np.linspace(x_param['range'][0], x_param['range'][1], x_param['points'])
        y_values = np.linspace(y_param['range'][0], y_param['range'][1], y_param['points'])
        X, Y = np.meshgrid(x_values, y_values)
        
        # 初始化流型矩阵
        regimes = np.zeros_like(X, dtype=object)
        
        # 计算每个点的流型
        for i in range(X.shape[0]):
            for j in range(X.shape[1]):
                # 设置参数
                params = fixed_params.copy()
                params[x_param['name']] = X[i, j]
                params[y_param['name']] = Y[i, j]
                
                # 确保alpha和velocity_r参数存在
                alpha = params.pop('alpha') if 'alpha' in params else 0.5
                velocity_r = params.pop('velocity_r') if 'velocity_r' in params else 0.5
                
                # 识别流型
                regimes[i, j] = self.identifier.identify_regime(alpha, velocity_r, **params)
        
        return X, Y, regimes
    
    def plot_map(self, x_param, y_param, **fixed_params):
        """
        绘制流型图
        
        Args:
            x_param: X轴参数配置 {'name': 参数名, 'range': [最小值, 最大值], 'points': 点数}
            y_param: Y轴参数配置 {'name': 参数名, 'range': [最小值, 最大值], 'points': 点数}
            **fixed_params: 固定参数值
            
        Returns:
            matplotlib.figure.Figure: 图形对象
        """
        try:
            import matplotlib.pyplot as plt
            from matplotlib.colors import ListedColormap
            
            # 生成流型数据
            X, Y, regimes = self.generate_map(x_param, y_param, **fixed_params)
            
            # 将流型转换为数值
            regime_types = list(set(regimes.flatten()))
            regime_map = {regime: i for i, regime in enumerate(regime_types)}
            regime_values = np.zeros_like(X, dtype=int)
            for i in range(X.shape[0]):
                for j in range(X.shape[1]):
                    regime_values[i, j] = regime_map[regimes[i, j]]
            
            # 创建颜色映射
            colors = ['blue', 'green', 'yellow', 'red', 'purple']
            cmap = ListedColormap(colors[:len(regime_types)])
            
            # 绘制流型图
            fig, ax = plt.subplots(figsize=(10, 8))
            contour = ax.contourf(X, Y, regime_values, levels=len(regime_types), cmap=cmap)
            
            # 添加标签和图例
            ax.set_xlabel(x_param['name'])
            ax.set_ylabel(y_param['name'])
            ax.set_title('流型图')
            
            # 创建图例
            legend_elements = [plt.Rectangle((0, 0), 1, 1, color=colors[i]) 
                              for i in range(len(regime_types))]
            ax.legend(legend_elements, regime_types, loc='best')
            
            # 添加颜色条
            plt.colorbar(contour, ax=ax, ticks=range(len(regime_types)))
            
            return fig
        except ImportError:
            print("需要安装matplotlib库来绘制流型图")
            return None


# 使用示例
if __name__ == "__main__":
    # 创建流型识别器
    identifier = FlowRegimeIdentifier()
    
    # 基本流型识别
    print("基本流型识别:")
    print(f"空泡率=0.2, 相对速度=0.3 m/s: {identifier.identify_regime(0.2, 0.3)}")
    print(f"空泡率=0.4, 相对速度=0.5 m/s: {identifier.identify_regime(0.4, 0.5)}")
    print(f"空泡率=0.6, 相对速度=0.8 m/s: {identifier.identify_regime(0.6, 0.8)}")
    print(f"空泡率=0.8, 相对速度=1.2 m/s: {identifier.identify_regime(0.8, 1.2)}")
    
    # 使用更多参数
    print("\n使用更多参数:")
    params = {
        'diameter': 0.02,      # 2cm
        'rho_l': 998.0,        # 水的密度
        'rho_g': 1.2,          # 空气的密度
        'mu_l': 1.0e-3,        # 水的粘度
        'mu_g': 1.8e-5,        # 空气的粘度
        'sigma': 0.073,        # 水-空气的表面张力
        'is_vertical': True    # 垂直流动
    }
    print(f"空泡率=0.4, 相对速度=0.5 m/s, 更多参数: {identifier.identify_regime(0.4, 0.5, **params)}")
    
    # 创建流型图
    flow_map = FlowRegimeMap(identifier)
    
    # 尝试绘制流型图
    try:
        x_param = {'name': 'alpha', 'range': [0.1, 0.9], 'points': 20}
        y_param = {'name': 'velocity_r', 'range': [0.1, 2.0], 'points': 20}
        fixed_params = {
            'diameter': 0.02,
            'rho_l': 998.0,
            'rho_g': 1.2,
            'mu_l': 1.0e-3,
            'mu_g': 1.8e-5,
            'sigma': 0.073,
            'is_vertical': True
        }
        fig = flow_map.plot_map(x_param, y_param, **fixed_params)
        if fig:
            fig.savefig('flow_regime_map.png')
            print("\n流型图已保存为 'flow_regime_map.png'")
    except Exception as e:
        print(f"\n绘制流型图失败: {e}")
